from dotenv import load_dotenv
from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate

# 确保你也设置了OpenAI的API密钥
import os

from langchain_openai import AzureChatOpenAI

load_dotenv()

llm_model = AzureChatOpenAI(
    # openai_api_key=
    # openai_api_base=os.getenv("AZURE_OPENAI_ENDPOINT"),
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    azure_deployment=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
    api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
    temperature=0.7
)

# 创建一个简单的链
prompt = ChatPromptTemplate.from_template("给我写一个关于{topic}的短诗")
chain = prompt | llm_model

# 运行链
result = chain.invoke({"topic": "秋天"})
print(result)
